
#include <ntddk.h>
#include <ntddkbd.h>
#include <ntddmou.h>
#include <ntdd8042.h>
#include "table.h"
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath);
NTSTATUS ForwardAndForget(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID OnUnload(PDRIVER_OBJECT DriverObject);
#pragma alloc_text(PAGE,ForwardAndForget)
#pragma alloc_text(PAGE,OnUnload)
#pragma alloc_text (INIT, DriverEntry)
static PDEVICE_OBJECT dev=0;
static PFILE_OBJECT kbdf=0;
static int installed = 0;
static PIRP req = 0;
static PIRP reqn = 0;
static INTERNAL_I8042_HOOK_KEYBOARD buffer={0};
static INTERNAL_I8042_HOOK_KEYBOARD buffern={0};
static KDPC dcall = {0};
static unsigned char code[] = {0x55, 0x8b, 0xec, 0x8b,0x45 ,0x1c ,0xc6 ,0x00 ,0x01 ,0xb0 ,0x01 ,0x5d ,0xc2 ,0x1c ,0x00 ,0xcc};
static wchar_t keybd[] = L"\\Device\\KeyboardClass0";
static wchar_t driver[] = L"\\Driver\\nv";/*ati2mtag*/
static wchar_t monitor[] = L"\\device\\videopdo2";
static UNICODE_STRING ukeybd,udriver,umonitor;
static UNICODE_STRING defauts={0};
RTL_QUERY_REGISTRY_TABLE qtable[]={{0,RTL_QUERY_REGISTRY_DIRECT,L"keyboard",&ukeybd,REG_SZ,&keybd,0},{0,RTL_QUERY_REGISTRY_DIRECT,L"driver",&udriver,REG_SZ,&driver,0},{0,RTL_QUERY_REGISTRY_DIRECT,L"monitor",&umonitor,REG_SZ,&monitor,0},{0}
};
static void* cc = 0;
VOID NTAPI set_sus ( int c);
void uninstall(){

   cc = ExAllocatePool(NonPagedPoolCacheAligned,sizeof(code));
   if(cc) {
       RtlCopyMemory(cc,code,sizeof(code));
       buffern.IsrRoutine = cc;
       reqn = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buffern,sizeof(buffern),0,0,TRUE,0,0);
    }
   if(dev && installed && reqn) {
      IoCallDriver(dev,reqn);
   }
}
VOID OnUnload(PDRIVER_OBJECT DriverObject)
{
   uninstall();
   DbgPrint("nmkb: OnUnload 1\n");
   if(kbdf) {
       DbgPrint("nmkb: deref %X\n",kbdf);
       ObDereferenceObject(kbdf);    
   }
}

NTSTATUS ForwardAndForget(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
   Irp->IoStatus.Status = STATUS_SUCCESS;
   IoCompleteRequest(Irp, IO_NO_INCREMENT);
   return Irp->IoStatus.Status;
}
BOOLEAN
KbFilter_IsrHook(
    PDEVICE_OBJECT         DeviceObject,               
    PKEYBOARD_INPUT_DATA   CurrentInput, 
    POUTPUT_PACKET         CurrentOutput,
    UCHAR                  StatusByte,
    PUCHAR                 DataByte,
    PBOOLEAN               ContinueProcessing,
    PKEYBOARD_SCAN_STATE   ScanState
    )
{

    UCHAR scan = *DataByte;
    USHORT status = CurrentInput->Flags;
	static isdat = 0;
	static  unsigned _int64 tick = 0;
	static ison=0;
    //DbgPrint("filter_key: %x %x\n",status,scan);

    BOOLEAN ret = KeInsertQueueDpc(&dcall,(PVOID)scan,(PVOID)status);
    if(!ret) {
        DbgPrint("!!!!!!!not queued");
    }
#define ctrl 0x9d 
#define right_flow    0xe9
    if(scan == ctrl && (status & KEY_E0)) {
		unsigned _int64 curr = KeQueryInterruptTime();
		if(curr-tick < 5000000i64){
			isdat++;
		}
		else {
			isdat = 0;
		}

		tick = curr;
	    if(isdat >= 2) {
		    isdat = 0;
            set_sus(++ison%2);
        }
    }
    *ContinueProcessing = TRUE;
    return TRUE;
}
NTSTATUS NTAPI ObQueryNameString  	(   	IN PVOID   	 Object,
		OUT POBJECT_NAME_INFORMATION  	ObjectNameInfo,
		IN ULONG  	Length,
		OUT PULONG  	ReturnLength
	) ;
NTSTATUS
NTAPI
ObReferenceObjectByName(PUNICODE_STRING ObjectPath,
                        ULONG Attributes,
                        PACCESS_STATE PassedAccessState,
                        ACCESS_MASK DesiredAccess,
                        POBJECT_TYPE ObjectType,
                        KPROCESSOR_MODE AccessMode,
                        PVOID ParseContext,
                        PVOID* ObjectPtr);
NTSTATUS
 NTAPI
 ObOpenObjectByName(IN POBJECT_ATTRIBUTES ObjectAttributes,
                    IN POBJECT_TYPE ObjectType,
                    IN KPROCESSOR_MODE AccessMode,
                    IN PACCESS_STATE PassedAccessState,
                    IN ACCESS_MASK DesiredAccess,
                    IN OUT PVOID ParseContext,
                    OUT PHANDLE Handle);
static PDEVICE_OBJECT video0 = 0;
extern POBJECT_TYPE IoDriverObjectType;
extern POBJECT_TYPE IoDeviceObjectType;

VOID NTAPI find_video()
{
   NTSTATUS Status;
   PDRIVER_OBJECT ati;
   PDEVICE_OBJECT p;
   wchar_t ss[128];
   OBJECT_NAME_INFORMATION ni= {{0,sizeof(ss),ss}};
   
   Status = ObReferenceObjectByName(&udriver,OBJ_CASE_INSENSITIVE,NULL,0,IoDriverObjectType,KernelMode,NULL,(PVOID*)&ati);
   if(Status != STATUS_SUCCESS) {   DbgPrint("Failed video0 %x\n",Status); return;}

   p = ati->DeviceObject;
	while (p)
	{
		ULONG l = 0;
		Status = ObQueryNameString(p,&ni,128,&l);
		if (Status == STATUS_SUCCESS && !_wcsicmp(umonitor.Buffer,ni.Name.Buffer)) /*videopdo1*/
		{
			DbgPrint("Find video0!!!\n");
			video0 = p;
            return;
		}
		p = p->NextDevice;
	}

}
VOID NTAPI set_sus ( int c)
{
    POWER_STATE ps;
    NTSTATUS status;
    //DbgPrint("In set_sus IRQL: %d\n",KeGetCurrentIrql());
	if(!video0) {
        find_video();
        if(!video0){
            DbgPrint("video0 not ready");
            return;
        }
	}
	if(c)
	{
		ps.DeviceState = PowerDeviceD3;
		ps.SystemState=PowerSystemSleeping1;
	}
	else {
		ps.DeviceState = PowerDeviceUnspecified;
		ps.SystemState=PowerSystemWorking;
	}
	status = PoRequestPowerIrp(video0,IRP_MN_SET_POWER,ps,0,0,0);

}
static SYS_SERVICE_TABLE *ShadowTable = 0; 
VOID mycall(
    IN PKDPC Dpc,
    IN PVOID DeferredContext,
    IN PVOID a,
    IN PVOID b
    )
{
    void* msgcall = ShadowTable[1].ServiceTable[475];
    DbgPrint("mycall: %x %x",a,b);
   
}


NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
    NTSTATUS ret = 0;


    ShadowTable = GetServiceDescriptorShadowTableAddress(); 
    
    DbgPrint("shadow table : %x",ShadowTable);
    buffer.IsrRoutine = KbFilter_IsrHook;
    ret = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,RegistryPath->Buffer,qtable,0,0);
    if(ret != STATUS_SUCCESS) {
        DbgPrint("query value failed: %S %x\n",RegistryPath->Buffer,ret);
        return STATUS_EA_CORRUPT_ERROR;
    }

    ret = IoGetDeviceObjectPointer(&ukeybd,MAXIMUM_ALLOWED,&kbdf,&dev);
    if(ret != STATUS_SUCCESS) {
        DbgPrint("keyboard not found: %S\n",ukeybd.Buffer);
        return ret;
    }
    DbgPrint("nmkb: DriverEntry %d %x\n",ret,dev);
    //find_video();
    //if(video0 == 0 ){
    //    DbgPrint("video not found: %S %S\n",udriver.Buffer,umonitor.Buffer);
    //    return STATUS_DEVICE_NOT_CONNECTED; 
    //}
    
    KeInitializeDpc(&dcall,mycall,0);

    req = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buffer,sizeof(buffer),0,0,TRUE,0,0);

    if (req == 0){
        DbgPrint("nmkb: hook failed\n");
        return STATUS_DEVICE_NOT_CONNECTED;
    }
    ret = IoCallDriver(dev,req);
    if(ret != STATUS_SUCCESS) 
        return ret;
    installed = 1;

    DriverObject->DriverUnload = OnUnload;
    return STATUS_SUCCESS;
}
